program program_psr
use madx_ptc_module
use pointer_lattice
implicit none

    interface
       subroutine build_PSR(PSR,MIS)
         use madx_ptc_module
         use pointer_lattice
         implicit none
         type(layout), target :: PSR
         logical(lp) mis
       end subroutine build_PSR
    end interface

type(layout), pointer:: psr
type(damap) id,m
type(real_8) ray(6)
real(dp) closed_orbit(6)
type(internal_state) state
type(normalform) normal
integer :: no1=0,np1=0
real(dp), allocatable :: average_floquet(:)
logical(lp) mis
integer, allocatable :: jexp(:)
integer i,j,d_of_f,k

call ptc_ini

psr=>m_u%start

Write(6,*) " small misalignments and rotations in lattice ? input t or f "
read(5,*) mis

call build_PSR(PSR,mis)

!!!  Phase space of Taylor maps will be (x,px,y,py)
state=only_4d0


 no1=1
 write(6,*) " order = ",no1     ! set to 2 because no parameters here

np1=0    !!!  No system parameters in the map

call init(state,no1,np1) ! PTC and FPP are properly initialized

call alloc(id,m)  
call alloc(ray)
call alloc(normal)

closed_orbit=0.d0 ! initial guess for closed orbit

call find_orbit(psr,closed_orbit,1,state,1.d-5)

write(6,'(a16,4(1x,g12.5))') " closed orbit = ",closed_orbit(1:4)

id=1;   ray=closed_orbit+id;   ! ray= closed orbit + identity map

call track(psr,ray,1,state) ! One turn map is computed via the ray

!!! The ray is really the result of a truncated power series algebra
!!! The map we are interested to compute is around the closed orbit: it is really part of a diffential algebra
m=ray   !   The ray is "officially" turned into a "damap"  (DA in Berz's honour)

normal = m  ! The map is normalised   

write(6,*)   " "
write(6,*)   " <z_i*z_j> = (A_i_1*A_j_1+A_i_2*A_j_2)*J_1+(A_i_3*A_j_3+A_i_4*A_j_4)*J_2+..."
write(6,*)   " "
allocate(jexp(c_%nd2))
allocate(average_floquet(c_%nd))
do i=1,c_%nd2
do j=i,c_%nd2
do d_of_f=1,c_%nd
 jexp=0
 jexp(2*d_of_f-1)=1
 average_floquet(d_of_f)=(normal%A_T%v(i).sub.jexp)*(normal%A_T%v(j).sub.jexp)
 jexp=0
 jexp(2*d_of_f)=1
 average_floquet(d_of_f)=(normal%A_T%v(i).sub.jexp)*(normal%A_T%v(j).sub.jexp)+ average_floquet(d_of_f)
 enddo
write(6,'(2x,a3,i1,a3,i1,a4,4(g18.11,a3,i1,a3))')"<z_", i,"z_",j,"> = ",( average_floquet(k)," J_",k," + ", & 
k=1,c_%nd-1),average_floquet(c_%nd)," J_",c_%nd
enddo
enddo

write(6,*) "   "
write(6,*) " hit return to terminate program "
write(6,*) "   "
pause 

deallocate(jexp)
deallocate(average_floquet)

call kill(id,m)  
call kill(ray)
call kill(normal)

end program program_psr


!=================================================================
subroutine  build_PSR(PSR,mis)
use madx_ptc_module
use pointer_lattice
implicit none

type(layout), target :: PSR

real(dp) :: ang, brho, kd, kf, Larc,SG(6),cut
type(fibre) :: b, d1, d2, qd, qf
type(layout) :: cell
logical(lp) mis
!-----------------------------------

call make_states(.false.)
exact_model = .true.
default = default + nocavity  
call update_states
madlength = .false.

ang = (twopi * 36.d0 / 360.d0)
Larc = 2.54948d0
brho = 1.2d0 * (Larc / ang)
call set_mad(brho = brho, method = 2, step = 10)
madkind2 = drift_kick_drift

kf =  2.72d0 / brho
kd = -1.92d0 / brho

d1 = drift("D1", 2.28646d0)
d2 = drift("D2", 0.45d0)
qf = quadrupole("QF", 0.5d0, kf)
qd = quadrupole("QD", 0.5d0, kd)
b  = rbend("B", Larc, ang)
cell = d1 + qd + d2 + b + d2 + qf + d1

PSR = 10 * cell
PSR = .ring.PSR

call survey(PSR)

if(mis) then
 sig=1.d-5
 cut=4.d0
 call MESS_UP_ALIGNMENT(PSR,SIG,cut)
endif
end subroutine build_PSR
